from turtle import pd

import torch

torch.manual_seed(1314)
from torch import nn
import torchvision.models as models
import torchvision.transforms as transforms
import cv2

import tqdm
import torch
import torch.utils.data
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import os
import pandas as pd


class FeatureExtractor(nn.Module):# 提取特征工具
    def __init__(self, submodule, extracted_layers):
            super(FeatureExtractor, self).__init__()
            self.submodule = submodule
            self.extracted_layers = extracted_layers

    def forward(self, x):
        outputs = []
        for name, module in self.submodule._modules.items():
            if name == "fc":
                x = x.view(x.size(0), -1)
            x = module(x)
            if name in self.extracted_layers:
                outputs.append(x)
        return outputs